import sys,os
from collections import defaultdict
import pyfaidx
import argparse




def parser_agpfile(agpfile, outagp):
    names = locals()
    dic = defaultdict(dict)
    with open(agpfile) as agp:
        for i in agp.readlines():
            tmp = i.strip().split("\t")
            ctg = tmp[5]
            if not ctg.startswith("h"):
                continue
            ctg_start = int(tmp[6])
            ctg_end = int(tmp[7])
            key = f"{ctg_start}-{ctg_end}"
            dic[ctg][key] = i
    #print(dic)
    agp.close()
    error_ctg = []
    for key in dic.keys():
        if len(dic[key].keys()) > 1:
            print(f"{key} has more than one AGP line:")
            error_ctg.append(key)
    
    
    num = 1
    ctg_list = []
    with open(agpfile) as agp, open(outagp, "w") as out:
        for s in agp.readlines():
            tmp = s.strip().split("\t")
            ctg = tmp[5]
            if not ctg.startswith("h"):
               out.write(s)
               continue
            # if ctg not in ctg_list:
            #     names[f"{ctg}_num"] = 1
            #     ctg_list.append(ctg)
            # else:
            #     names[f"{ctg}_num"] = names[f"{ctg}_num"] + 1
            # if ctg in error_ctg:
    
            #     outline = s.replace(ctg, f"{ctg}_{num_ctg}")
            #     out.write(outline)
            if ctg in error_ctg:
                start = tmp[6]
                end = tmp[7]
                length = int(end) - int(start) + 1
                newstart = 1
                newend = length
                newctgname = f"{ctg}:{start}-{end}"
                outline = "\t".join(tmp[0:5]) + "\t" + newctgname + "\t" + str(newstart) + "\t" + str(newend) + "\t" + tmp[8] + "\n"
                out.write(outline)

            else:
                out.write(s)






        
def parse_fasta(infasta, outfasta, agpfile):
    """
    parse fasta file
    """
    # dic = {}
    # with open(infasta) as fasta, open(outfasta, "w") as out:
    #     name = "name"
    #     seq = "seq"
    #     for s in fasta:
    #         if s.startswith(">"):
    #             dic[name] = seq
    #             name = s.strip()[1:]
    #             seq = ""
    #         else:
    #             seq = seq + s.strip()
    #     dic[name] = seq
    #     for line in open(agpfile):
    #         agptmp = line.strip().split("\t")
    fasta = pyfaidx.Fasta(infasta)
    with open(agpfile) as agp, open(outfasta, "w") as out:
        for line in agp:
            tmp = line.strip().split("\t")
            #print(line)
            ctg = tmp[5]
            if not ctg.startswith("h"):
                continue
            if ":" in ctg:
                ctgraw = ctg.split(":")[0]
                start, end = [int(i) for i in ctg.split(":")[1].split("-")]
                try:
                    seq = str(fasta[ctgraw][start - 1:end])
                except:
                    seq = str(fasta[ctgraw][start - 1:])
                out.write(f">{ctg}\n")
                for sq in range(0, len(seq), 60):
                    out.write(seq[sq:sq + 60] + "\n")
            else:
                seq = str(fasta[ctg][0:])
                out.write(f">{ctg}\n")
                for sq in range(0, len(seq), 60):
                    out.write(seq[sq:sq + 60] + "\n")
                

def main():
    parser = argparse.ArgumentParser(description="correct fasta file by agpfile")
    parser.add_argument("-i", "--input", help="input agp file", required=True, dest="inagp")
    parser.add_argument("-o", "--output", help="output fasta file", required=True, dest="output")
    parser.add_argument("-f", "--fasta", help="input fasta file", required=True, dest="fasta")
    parser.add_argument("-a", "--agp", help="the output new agpfile", required=True, dest="outagp")
    args = parser.parse_args()
    
    inagp = args.inagp
    outagp = args.outagp
    fasta = args.fasta
    outfasta = args.output
    parser_agpfile(inagp, outagp)     
    parse_fasta(fasta, outfasta, outagp)

if __name__ == "__main__":
    main()
  
        
        
        
        